# this is for primer design for tiling larger targets
# criteria is set by the batcher that initiates this spript
# it is something like mean+2SD
# for now use it for intermediate and long sequences but for very short sequences it is not that efficient since it has to generate primer lists many times, rather than only once. 


import os
import sys
sys.path.append(os.path.join(os.getcwd(), "modules"))
import cPickle
import random
import pysam

from itertools import product
# import custom functions
from allfunctions import revcomp, Tm, dimer, GC
from modefunctions import find_exon, check_snp

def pck_large(outputdir, filename, pnumber, header, fullseq, runDB, pathDB, Random):
    samp = 200

    designmode = runDB["designmode"]
    POI = runDB["POI"]

    ############ DEFINE FUNCTIONS ##############
    # this is the picker function 
    # the function input is a list
    if designmode == 'PCR':
        from modefunctions import pcr_picker as picker
        #dna=250.0; Na=50.0; K=0; Tris=0; Mg=1.5; dNTP=0.2; DMSO=0

        Arm_max = "na" # total length of forward and reverse arms allowed
        REdict = "na"
        PriF = "na" # Forward MIPamp primer site
        Middle = "na" # The backbone
        PriR = "na" # Reverse MIPamp primer site
        Oligo_len = "na" # Lenght of the final product for ordering

        Remove = False
        if runDB["repeats"] == "ignore":
            from modefunctions import pcr_score_pair_IR as score_pair
        elif runDB["repeats"] == "score" :
            from modefunctions import pcr_score_pair_S as score_pair
        elif runDB["repeats"] == "remove":
            from modefunctions import pcr_score_pair_IR as score_pair
            Remove = True

    elif designmode == 'MIP':
        from modefunctions import mip_picker as picker
        #dna=10; Na=5; K=27; Tris=22.9; Mg=10; dNTP=0.3; DMSO=0
        # use values from ampligase buffer, D1/20 Ampligase storage buffer and, D1/50 Phusion storage buffer

        Arm_max = runDB['Arm_max'] # total length of forward and reverse arms allowed
        REdict = runDB['REdict'] 
        PriF = runDB['PriF'].lower() # Forward MIPamp primer site
        Middle = runDB['Middle'].lower() # The backbone
        PriR = runDB['PriR'].lower() # Reverse MIPamp primer site
        Oligo_len = runDB['Oligo_len'] # Lenght of the final product for ordering

        Remove = False
        if runDB["repeats"] == "ignore":
            from modefunctions import score_mip_IR as score_pair
        elif runDB["repeats"] == "score":
            from modefunctions import score_mip_S as score_pair
        elif runDB["repeats"] == "remove":
            from modefunctions import score_mip_IR as score_pair
            Remove = True

    #### END OF FUNCTIONS ####

    ##########################################
    # ENTER RUN SPECIFIC CRITERIA
    MinGC, MaxGC = runDB['MinGC'], runDB['MaxGC']
    MinTm, MaxTm = runDB['MinTm'], runDB['MaxTm']
    MinLen, MaxLen = runDB['MinLen'], runDB['MaxLen']
    Epos = runDB['Epos'] # relative position of the exon in the target (each end)
    Smin, Smax = runDB['Smin'], runDB['Smax'] 
    Olap = runDB['Olap'] # tiling overlap
    pntcutoff = runDB['pntcutoff']
    dna = runDB["dna"]
    Na = runDB["Na"]
    K = runDB["K"]
    Tris = runDB["Tris"]
    Mg = runDB["Mg"]
    dNTP = runDB["dNTP"]
    DMSO = runDB["DMSO"]
    tiling = runDB["tiling"]
    #########################################


    # all the output file data will be collected in a dict to be PIPE_OUT
    output_data = {"display":{pnumber: list()}, "picked":list()}

    tcount = 1
    #pnumber= int(filename.split("-")[1]) #for primer naming (K)
    failed = "NA" # count number of primer failed by SNP

    TTemp2 = list() # the master picked list
    len_TTemp2 = 0
    old_pick = list()
    pick_pass = 0
    Picked_code = dict()
    loopbreak = False
    bases = set(['a', 'c', 'g', 't'])    

    # i counter for region, k for count F, t for R 
    i = 0
    k = 1
    t = 1
    Pi = 1 # this is to mark each pair with a unique name

    genome, Chr, fullTstart, fullTstop, Offset = header

    repeatDB = pathDB["repeatDB"]
    exonfile = pysam.Tabixfile(os.path.join(pathDB["exonDB"], "%s.tabix.gz" %Chr))
    snpfile = pysam.Tabixfile(os.path.join(pathDB["snpDB"], "%s.tabix.gz" %Chr))

    # now there is the real header from the bed file
    # and then there is the adjusted one
    # at each stem you need to make sure you use the correct one
    lenTarget= fullTstop-fullTstart
    # Epos is not only Epos but also Offset too
    Epos += Offset
    Estart = fullTstart + Epos
    Estop = fullTstop - Epos
    target = '%s:%s-%s' %(Chr, Estart, Estop)

    scale = max((lenTarget)/100, 1)
    # to be used in the display file for scaling
    # get first part of the seq
    seqstart = 0
    seqstop = Epos+Smax

    # generate a display output file
    output_data["display"][pnumber].append('\n%s\t%s:%s-%s\n' %tuple(header[0:4]))
    output_data["display"][pnumber].append('\nRegion %s Scale is %d\n' %(pnumber, scale))

    if not POI:
        output_data["display"][pnumber].append('%s%s%s\n' %('-'*(Epos/scale), '|'*((Estop-Estart)/scale), '-'*(Epos/scale)))

    else:
        output_data["display"][pnumber].append('%s*%s\n' %('|'*(len(fullseq)/(2*scale)), '|'*(len(fullseq)/(2*scale))))


    RstartN = fullTstart
    FstopN = fullTstart
    Lsize = Smin
    last_run = False
    pick_reverse = False
    while 1:
        seq = fullseq[seqstart:seqstop]
        R= list() # reverse primers list
        F= list() # forwrd primers list
        # there is a break statement at the end that will stop the loop
        Tstart = fullTstart + seqstart #start coordinates of the subsequence
        Tstop = min((Tstart + len(seq)-1), fullTstop) #stop coordinates of the subsequence
        if len(seq) < Smin:
            pick_reverse= None
            # you have to pick a reverse primer
            seq= fullseq[-(Smax+Epos):]
            Tstart= fullTstop - Smax + 1
            Tstop= fullTstop

        if not POI:
            # define seqR
            seqR= revcomp(seq)

        else:
            # if poi then F has to be on the left, R has to be on the right
            seqR= revcomp(fullseq[Epos + 2 * MaxLen:])
            seq= fullseq[:Epos - 2 * MaxLen + 1]
            # so do not waste time designing others

        l= -1
        m= MinLen-2
        n= MaxLen
        site= len(seq)
        while l < site-(n+1):
            l+=1
            i= 0
            m+=1
            M= m+1
            while i < (n-MinLen+1): 
                Fprimer= seq[l:M]
                len_Fprimer= len(Fprimer)
                Fstart= Tstart+l
                Fstop= Tstart+M-1
                if Fstop <= Tstop-Smin and sum([Fprimer.count(p) for p in bases]) == len_Fprimer:
                    # the primer should not start near the Tstop, unrelevant
                    # the primer should not have any other base then a,c,g,t
                    tm= Tm(Fprimer, dna, Na, K, Tris, Mg, dNTP, DMSO)
                    gc= GC(Fprimer)
                    if MinTm <= tm <= MaxTm and MinGC <= gc <= MaxGC:
                        if check_snp(snpfile, Chr, Fstart, Fstop, "F") == 'PASS':
                            if dimer(Fprimer, Fprimer, pntcutoff=pntcutoff) == 'PASS':
                                name= '%sF%s' %(pnumber, k)
                                F.append([name, Fprimer, tm, gc, len_Fprimer, target, Chr, Fstart, Fstop, Epos-M])
                                k+=1
                Rprimer= seqR[l:M]
                len_Rprimer= len(Rprimer)
                Rstart= Tstop-M+1
                Rstop= Tstop-l
                if Rstart >= Tstart+Smin and sum([Rprimer.count(p) for p in bases]) == len_Rprimer:
                    # the primer should not start near the Tstop, unrelevant
                    # the primer should not have any other base then a,c,g,t
                    tm= Tm(Rprimer, dna, Na, K, Tris, Mg, dNTP, DMSO)
                    gc= GC(Rprimer)
                    if MinTm <= tm <= MaxTm and MinGC <= gc <= MaxGC:
                        if check_snp(snpfile, Chr, Rstart, Rstop, "R") == 'PASS':
                            if dimer(Rprimer, Rprimer, pntcutoff=pntcutoff) == 'PASS':
                                name= '%sR%s' %(pnumber, t)
                                R.append([name, Rprimer, tm, gc, len_Rprimer, target, Chr, Rstart, Rstop, Epos-M])
                                t+=1
                M+=1
                i+=1


        # TTemp= [] # temporary for primer pairs
        # generate each pair and score each pair
        # random select at most 100 F and R, pair them in a heapq
        if Random:
            if designmode == 'PCR':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(random.sample(F, min(len(F), samp)), random.sample(R, min(len(R), samp))), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax]

            elif designmode == 'MIP':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(random.sample(F, min(len(F), samp)), random.sample(R, min(len(R), samp))), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax and pair[0][4] + pair[1][4] <= Arm_max]

        else:
            if designmode == 'PCR':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(F, R), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax]

            elif designmode == 'MIP':
                TTemp= [[0] + pair[0] + pair[1] + [(pair[1][8]-pair[0][7]+1), '%sP%s' %(pnumber, Pi)] for Pi, pair in enumerate(product(F, R), Pi) if Smin <= (pair[1][8]-pair[0][7]+1) <= Smax and pair[0][4] + pair[1][4] <= Arm_max]

        # advance the pair number (pi)
        # so that after ever sub-target they are still unique 


        # THE TILER FUNCTION
        # here goes the picking function for tiling
        picking= 'FAIL'
        old_pick= []

        len_TTemp2= len(TTemp2)
        if len_TTemp2 == 0 or len_TTemp2 % 2 == 0:
            # even
            even= True
        elif len_TTemp2 % 2 != 0:
            # odd
            even= False

        #print RstartN

        # NOW
        # ARE YOU TILING BASES AT X-FOLD
        # OR JUST TILING THE REGIONS
        if runDB["tiling"]:
            if Estop >= RstartN > fullTstart:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > RstartN + MaxLen]
                # have a fail safe
                if not len(TTemp1):
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < Li[9] < RstartN and Li[18] > RstartN + MaxLen]

            elif RstartN > Estop:
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if FstopN < Li[8] < RstartN - Lsize*(1-1.0/tiling) and Li[18] > Estop]
                # have a fail safe
                if not len(TTemp1):
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] < RstartN and Li[18] > Estop]
                    """
                    TTemp1= []
                    for Li in TTemp:
                        if FstopN < Li[8] < Li[9] < Li[18] < Li[19] < RstartN:
                            TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                        elif Li[9] < FstopN - MaxLen < RstartN + MaxLen < Li[18]:
                            TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                        elif Li[9] < FstopN - MaxLen < Li[18] < Li[19] < RstartN:
                            TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                        elif FstopN < Li[8] < Li[9] < RstartN + MaxLen < Li[18]:
                            TTemp1.append([score_pair(Li, repeatDB, Smax, even)] + Li[1:])
                    if not len(TTemp1):                    
                        TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] < RstartN and Li[18] > Estop]
                        """

                last_run= True

            else:
                # this ensure that first tiling primer is picked
                # within the 5' buffer region outside the target exon
                # always try to cover the region with one product
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estart]


            #print len(TTemp1), 
            # let start picking some primers now
            if len(TTemp1) > 0:
                Picked_code, thepair= picker(TTemp1, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, repeatDB, Picked_code, runDB['gfserver'], Remove)

                if thepair != 'FAIL':
                    thepair.extend([find_exon(exonfile, Chr, thepair[8], thepair[19])])

                    TTemp2.append(thepair)
                    len_TTemp2= len(TTemp2)
                    picking= 'PASS'

                    # adjust RstartN
                    RstartN= TTemp2[-1][18]
                    Lsize= TTemp2[-1][21]
                    FstopN= TTemp2[-1][9]
                    #this is the Reverse start position of the last pair
                    #this is the Reverse stop position of the last pair 

            if last_run:
                break

            # now its time to go back to the top to move the region
            # in order to determine where the new region will start
            if picking == 'FAIL':
                RstartN += 5
                FstopN += 5

            # seqstart is the relative start position on the full seq
            seqstart= FstopN - fullTstart
            # seqstop is the relative stop position
            seqstop= seqstart + Smax + Olap


        else:
            if RstartN > Tstart:
                # any other case
                if fullTstop - RstartN + Olap <= Smin:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[18] > Estop]
                    if not len(TTemp1):
                        # you are miserable, just give me something, work with me
                        TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[8] > RstartN - Smin + 2*MaxLen]
                    pick_reverse == None

                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if RstartN-Olap <= Li[9] < RstartN and Li[18] > RstartN + min(Olap+2*MaxLen, Smax)]

                """
                rule 1:
                the new pairs Fstop cannot be smaller than the Rstop-Olap of the old pick
                but has to be smaller than Rstop of that primer

                rule 2:
                the new pairs Rstart has to be futher away from the Rstart of the old pick
                by Olap+2*MaxLen 
                (one for the offset from Rstop-Rstart, other from Li[9]- Li[8]), 
                so that the new Fstart won't overlap with the previous pairs
                Rstop
                """

            elif POI:
                # if it is a point-of-interest pick that is all you need
                # one good pick
                TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp]

            else:
                # this ensure that first tiling primer is picked
                # within the 5' buffer region outside the target exon
                # always try to cover the region with one product
                if fullTstop - fullTstart <= Smax:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estop]
                else:
                    TTemp1= [[score_pair(Li, repeatDB, Smax, even)] + Li[1:] for Li in TTemp if Li[9] < Estart and Li[18] > Estart]


            #print len(TTemp1)
            #now pick
            if len(TTemp1) > 0:
                Picked_code, thepair= picker(TTemp1, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, repeatDB, Picked_code, runDB['gfserver'], Remove)

                if thepair != 'FAIL':
                    thepair.extend([find_exon(exonfile, Chr, thepair[8], thepair[19])])

                    TTemp2.append(thepair)
                    len_TTemp2= len(TTemp2)
                    picking= 'PASS'
                    pick_pass += 1
                    if pick_pass == 2:
                        # you need two consequtive picks before clearing
                        old_pick= []
                        pick_pass= 0 
                    # adjust RstartN
                    RstartN= TTemp2[-1][18]
                    #this is the Reverse start position of the last pair
                    #this is the Reverse stop position of the last pair 
                    if RstartN > Estop:
                        break

            if POI or pick_reverse == None:
                # for POI
                # you get one shot whether it is pass or fail does not matter
                # there is nothing you can do to make it pass after fail
                break

            # now its time to go back to the top to move the region
            # in order to determine where the new region will start
            if picking == 'FAIL':
                if pick_pass == 1 and len(old_pick) == 1:
                    # you failed before and picked a new one but you cannot fool me
                    # you failed yet again
                    # don't forget to remove the new pick
                    # but which one to keep, they both are bad
                    try:
                        TTemp2.pop()
                    except IndexError:
                        # if TTemp2 is empty now
                        pass
                    TTemp2.append(old_pick.pop())
                    # you tried everything you can but it does not work
                    # fix len(TTemp2) so there won't be any endless loops going back
                    len_TTemp2= len(TTemp2)
                    # the old pick now determines the RstartN
                    RstartN= TTemp2[-1][18]
                    FstopN= TTemp2[-1][9]
                    # just a difficult region, move along
                    RstartN+= 5
                    FstopN+= 5

                else:
                    # you failed straight up
                    if len(TTemp2) > len_TTemp2:
                        # no worries, do over, try a new pick
                        if len(old_pick) == 0:
                            old_pick.append(TTemp2.pop())
                            if len(TTemp2) > 0:
                                RstartN= TTemp2[-1][18]
                                FstopN= TTemp2[-1][9]
                            else:
                                RstartN= Tstart
                                FstopN= Tstart
                        else:
                            # you failed even after you tried a new pick
                            TTemp2.append(old_pick.pop())
                            # the old pick now determines the RstartN
                            RstartN= TTemp2[-1][18]
                            FstopN= TTemp2[-1][9]
                            # it could not inititate primer picking
                            RstartN+= 5
                            FstopN+= 5
                            # use plus one in case minlen= maxlen
                            """
                            IMPORTANT
                            the two if and elif statements may counter each other
                            so make sure that the net effect is going forward
                            so that they don't end up in a loop
                            thus go back MinLen, go forward MaxLen
                            """
                    else:
                        # it could not inititate primer picking
                        RstartN+= 5
                        FstopN+= 5
                        # use plus one in case minlen= maxlen


            if RstartN > fullTstop or RstartN < fullTstart:
                break

            # seqstart is the relative start position on the full seq
            seqstart= RstartN - Olap - fullTstart - 1
            # seqstop is the relative stop position
            seqstop= seqstart + Smax + Olap + 1
            """
            if seqstop > fullTstop - fullTstart:
                if loopbreak:
                    break
                else:
                    loopbreak = True
            #print seqstart, seqstop
            """
        
        #END OF THE TILER FUNCTION
        if pick_reverse == None and len(TTemp2) > 1:
            if len(TTemp2) > 2:
                rstartn= TTemp2[-3][18]
            else:
                rstartn= Estart

            # make sure that the last two primers are not reduntant.
            # pick the one that cover the most seq
            l1= TTemp2.pop() # the reverse pick
            l2= TTemp2.pop() # last forward pick
            keep= False
            for each in [l1, l2]:
                if each[9] < rstartn < Estop < each[18]:
                    # who ever that is add it back
                    TTemp2.append(each)
                    keep= True
                    break
            # the previous elimination did not work, try nearest neighbour
            if not keep:
                if abs(l1[8]-l2[8]) <=10 or abs(l1[19]-l2[19]) <= 10:
                    # do not use valuble oligo space for the gain of a few extra bases
                    # pick the one that has the most coverage
                    cdb={}
                    for each in [l1,l2]:
                        covg= each[18]-each[9]
                        # this is unique covg
                        # correct for non-target covg
                        if each[10] > 0:
                            covg-= each[10]
                        if each[20] > 0:
                            covg-= each[20]
                        cdb[covg]= each
                    TTemp2.append(cdb[max(cdb.keys())])
                else:
                    # if nothing else but both back
                    TTemp2.extend([l2, l1])


    # This is common to both tiling bases tiling regions
    for line in TTemp2:
        output_data["display"][pnumber].append('%s%s%s\n' %(' '*((line[8]-fullTstart)/scale), '-'*((line[19]-line[8])/scale), ' '*((fullTstop-line[19])/scale)))
        output_data["picked"].append(line)
    output_data["display"][pnumber].append("\n")

    
    """
    with open(os.path.join(outputdir, '%s.pkl' %filename), "w") as output:
        cPickle.dump(output_data, output)
        """
    return output_data
